#!/usr/bin/env python 
# -*- coding:utf-8 -*-


from cmd import Cmd
import localapi
import traceback
import os,sys
from libs.utils.printer import *
class Cli(Cmd):
    def __init__(self):
        os.system("cls")
        Cmd.__init__(self)
        banner = '''
         ____   ___   ____ 
        |  _ \ / _ \ / ___|
        | |_) | | | | |    
        |  __/| |_| | |___ 
        |_|    \___/ \____|{}
        '''.format("")
        self.intro = banner
        self.prompt = 'poc>> '
    def do_help(self,param):
        param = param.strip()
        if param == "set":
            localapi.help_set()
        else:
            localapi.info_pocs()
    def do_info(self,param:str):
        param = param.strip()
        if param == "pocs":
            localapi.info_pocs()
        elif param.find(".") > 0:
            poc,model= param.split('.')
            if poc == "poc":
                localapi.info_poc(model)
        else:
            warn("命令输入错误!")
    def do_set(self,param):
        param = param.strip()
        order,argvs = param.split(" ")
        if order == "url":
            localapi.set_url(argvs)
        elif order == "headers":
            if argvs.find("mitm.") == 0:
                sysName = ""
                rlist = argvs.split(".")
                del rlist[0]
                sysName = ".".join(rlist)
                localapi.set_headers_mitm(sysName)
            else:
                localapi.set_headers(argvs)
        elif order == "thread":
            localapi.set_threads(int(argvs))
        elif order == "timeout":
            localapi.set_timeout(int(argvs))
        else:
            warn("命令错误!")
    def do_check(self,param):
        localapi.check_set()
    def do_exec(self,param:str):
        param = param.strip()
        order,argvs = param.split(" ")
        if order == "pocs":
            localapi.exec_pocs(argvs)
        elif order == "poc":
            localapi.exec_poc(argvs)
        else:
            warn("命令错误!")
    def do_exit(self,param):
        warn("退出!")
        sys.exit()
    def default(self,param):
        print("命令不存在!")

if __name__ == "__main__":
    try:
        cli = Cli()
        cli.cmdloop()
    except Exception as e:
        #warn(e)
        print(e)
        traceback.print_exc()